#!/usr/bin/env python
# encoding: utf-8
"""
test.py

Created by Rob Hetland on 2008-01-18.
Copyright (c) 2008 Texas A&M Univsersity. All rights reserved.
"""


def test_gridgen():
    xbry = array([10.,  5.,  5., 0.5,  0.,  5.,  5., 10.])
    ybry = array([10., 10.,  6., 7.5,  6.,  4.,  0.,  0.])
    beta = array([ 1.,  1., -1., 1. ,  1., -1.,  1.,  1.])
    
    def focus(x, y, xo=0.55, yo=0.45):
        xf = tan((x - xo)*2.0)
        yf = tan((y - yo)*2.0)
        xf -= xf.min()
        xf /= xf.max()
        yf -= yf.min()
        yf /= yf.max()
        return xf, yf
    
    # Run Gridgen
    grid = gridgen(xbry, ybry, beta, (36, 36), focus=focus, ul_idx=1)
    
    # Islands for masking
    xmask = array([8., 7., 6., 7.])
    ymask = array([4., 5., 4., 3.])
    grid.maskpoly(zip(xmask, ymask))
    
    return grid

def test_make_grid():
    yv, xv = mgrid[0:20, 0:20:0.5]
    grd = Grid(x=xv, y=yv)
    
    def rot2d(x, y, ang):
        'rotate vectors by geometric angle'
        xr = x*cos(ang) - y*sin(ang)
        yr = x*sin(ang) + y*cos(ang)
        return xr, yr
    
    xvr, yvr = rot2d(xv, yv, pi/4)
    
    grd_r = Grid(x=xvr, y=yvr)
    
    print 'pn match? ', allclose(grd.pn, grd_r.pn)
    print 'pm match? ', allclose(grd.pn, grd_r.pn)
    
    dx = 1.0/grd_r.pm
    print 'dx min, max, and mean = ', dx.min(), dx.max(), dx.mean()
    dy = 1.0/grd_r.pn
    print 'dy min, max, and mean = ', dy.min(), dy.max(), dy.mean()
    
    latv, lonv = mgrid[15.0:32.0, -100.0:-80.0]
    lonv[0:5, 0:5] = nan
    lonv=ma.masked_where(isnan(lonv), lonv)
    grd_geo = Grid(lon=lonv, lat=latv)
    dx = 1.0/grd_geo.pm
    print 'dx min, max, and mean = ', dx.min(), dx.max(), dx.mean()
    dy = 1.0/grd_geo.pn
    print 'dy min, max, and mean = ', dy.min(), dy.max(), dy.mean()
    print grd_geo.f

def test_masking():
    xi, yi = 20*random.rand(2, 100)
    hi = exp( -(xi-10.)**2/10.0 -(yi-10.)**2/5.0)
    h = grid.extrap_xy_to_grid(xi, yi, hi)
    pl.contour(grid.x_rho, grid.y_rho, h)
    print h
    pl.show()

def test_make_masked_grid():
    yv, xv = mgrid[0:20, 0:20:0.5]
    mask = ones(xv.shape, 'd')
    mask[0:10, 0:20] = 0.0
    xv = ma.masked_where(mask==0, xv)
    yv = ma.masked_where(mask==0, yv)
    return make_cart_grid(xv, yv)

def test_grid_3d():
    yv, xv = mgrid[0:10, 0:10:0.5]
    grd = Grid(x=xv, y=yv)
    grd.h = 11.0 + 19*random.rand(*grd.x_rho.shape)
    grd.theta_s = 5.0
    grd.theta_b = 1.0
    grd.hc = 5.0
    grd.N = 20
    print 'z_w =', grd.z_w[:,5,5]
    print 'h = ', grd.h[5,5]

def test_nc_grid():
    grd = nc_grid('/Users/rob/Projects/Merrimack/Grid/merrimack_large_grd.nc')
    print grd.__dict__.keys()
    
    grd = nc_grid('/Users/rob/Models/roms/roms-3.0/ocean_his.nc')
    print grd.__dict__.keys()

def test_write_roms_grid():
    """test write_roms_grid method of Grid class"""
    y, x = mgrid[0:1:100j, 0:1:100j]
    grd = Grid(x, y)
    grd.f = 1e-4
    grd.h = 10.0
    grd.write_roms_grid('cart_test.nc')
    print ' ### wrote cart_test.nc'
    
    lat, lon = mgrid[43:45:100j, -68:-70:100j]
    grdg = Grid(lon=lon, lat=lat)
    grdg.h = 10.
    grdg.f = 1.0e-4
    grdg.write_roms_grid('geo_test.nc')
    print ' ### wrote geo_test.nc'
    

def test_rho_to_vert():
    yv, xv = mgrid[0:20, 0:20:0.5]
    def rot2d(x, y, ang):
        'rotate vectors by geometric angle'
        xr = x*cos(ang) - y*sin(ang)
        yr = x*sin(ang) + y*cos(ang)
        return xr, yr
    
    print 'Verticies calculated from rho points: '
    for ang in arange(0, 2*pi, pi/8):
        xvr, yvr = rot2d(xv, yv, pi/4)
        grd = Grid(x=xvr, y=yvr)
        x, y = rho_to_vert(grd.x_rho, grd.y_rho, grd.pm, grd.pn, grd.angle)
        print '    All values close for rotation of %f?  %s' % \
              (ang, allclose(x, xvr) and allclose(y, yvr))
    
    # pl.plot(xvr, yvr, '-k')
    # pl.plot(xvr.T, yvr.T, '-k')
    # pl.plot(x, y, '-r')
    # pl.plot(x.T, y.T, '-r')
    # pl.show()

